/*
 * SSH Assistant - Helper for SSH key, agent, and tunnel management.
 *
 * Copyright (c) 2005-2006 Justin Mecham <justin@aspect.net>
 * Copyright (c) 2003 Michael K Link <mlink@synthemesc.com>
 *
 * This program is free software; you can redistribute it and/or modify
 * it under the terms of the GNU General Public License as published by
 * the Free Software Foundation; either version 2 of the License, or
 * (at your option) any later version.
 *
 * This program is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 * GNU General Public License for more details.
 *
 * You should have received a copy of the GNU General Public License
 * along with this program; if not, write to the Free Software
 * Foundation, Inc., 59 Temple Place, Suite 330, Boston, MA 02111-1307 USA
 */

#import <Cocoa/Cocoa.h>

#include <stdlib.h>
#include <sys/types.h>
#include <pwd.h>
#include <unistd.h>
#include <signal.h>

const char * logname = NULL;
pid_t ssh_agent_pid  = 0;
char * ssh_auth_sock = NULL;

static void setupLoginItem()
{
    
    NSLog(@"Setting up SSH Assistant Login Item");

    NSAutoreleasePool   * pool     = [[NSAutoreleasePool alloc] init];
    NSMutableDictionary * myDict   = [[NSMutableDictionary alloc] init];
    NSUserDefaults      * defaults = [[NSUserDefaults alloc] init];
    NSMutableArray      * loginItems;
    
    [defaults addSuiteNamed:@"loginwindow"];
    
    [myDict setObject:[NSNumber numberWithBool:NO] forKey:@"Hide"];
    [myDict setObject:[[NSBundle mainBundle] bundlePath] forKey:@"Path"];
    
    loginItems = [[NSMutableArray arrayWithArray:[[defaults persistentDomainForName:@"loginwindow"] objectForKey:@"AutoLaunchedApplicationDictionary"]] retain];
    
    [loginItems removeObject:myDict];
    [loginItems addObject:myDict];
    
    [defaults removeObjectForKey:@"AutoLaunchedApplicationDictionary"];
    [defaults setObject:loginItems
                 forKey:@"AutoLaunchedApplicationDictionary"];
    [defaults synchronize];
    
    [loginItems autorelease];
    [pool release];
    
}

static void extractData(NSData *readData)
{
	NSString * readString = [[[NSString alloc] initWithData:readData encoding:NSASCIIStringEncoding] autorelease];
	NSArray  * readStringComponents = [readString componentsSeparatedByString:@"\n"];
	
	for (unsigned int i = 0; i < [readStringComponents count]; i++)
    {

		NSArray * environmentComponents    = [[readStringComponents objectAtIndex:i] componentsSeparatedByString:@";"];
		NSArray * environmentSubComponents = [[environmentComponents objectAtIndex:0] componentsSeparatedByString:@"="];
		
		if ([environmentSubComponents count] > 0)
        {
			if ([[environmentSubComponents objectAtIndex:0] isEqualToString:@"SSH_AGENT_PID"])
            {
				ssh_agent_pid = atoi([[environmentSubComponents objectAtIndex:1] cString]);
			}
			else if ([[environmentSubComponents objectAtIndex:0] isEqualToString:@"SSH_AUTH_SOCK"])
            {
				ssh_auth_sock = (char *) [[environmentSubComponents objectAtIndex:1] cString];
			}
		}
	}
}

static int login_hook(int argc, const char *argv[], struct passwd *pwd)
{

	logname = argv[1];
	
	if ((setgid(pwd->pw_gid) == -1) || (setuid(pwd->pw_uid) == -1))
    {
        NSLog(@"setgid() failed. %s ?= %d", logname, pwd->pw_gid);
		exit(EXIT_FAILURE);
	}

    NSLog(@"Hello");
	NSUserDefaults      * defaults = [NSUserDefaults standardUserDefaults];
//	NSMutableDictionary * loginDefaults;
//	NSMutableArray      * AutoLaunchedApplicationDictionary;
//	NSDictionary        * sshAddDictionary = [NSDictionary dictionaryWithObjectsAndKeys:
//											   [NSNumber numberWithBool:NO], @"Hide",
//											   [[NSBundle mainBundle] bundlePath], @"Path",
//											   nil
//                                             ];
//
//	if (!(loginDefaults = [defaults persistentDomainForName:@"loginwindow"]))
//    {
//		[defaults setPersistentDomain:
//			[NSDictionary dictionaryWithObject:
//				[NSArray arrayWithObject:sshAddDictionary]
//				forKey:@"AutoLaunchedApplicationDictionary"
//			] 
//			
//			forName:@"loginwindow"
//		];
//	}
//	else
//    {
//		loginDefaults = [[defaults persistentDomainForName:@"loginwindow"] mutableCopy];
//		[loginDefaults autorelease];
//		
//		if (!(AutoLaunchedApplicationDictionary = [loginDefaults objectForKey:@"AutoLaunchedApplicationDictionary"]))
//        {
//			AutoLaunchedApplicationDictionary = [NSMutableArray arrayWithObject:sshAddDictionary];
//		}
//		else
//        {
//			AutoLaunchedApplicationDictionary = [[loginDefaults objectForKey:@"AutoLaunchedApplicationDictionary"] mutableCopy];
//			[AutoLaunchedApplicationDictionary autorelease];
//			for (unsigned int i = 0; i < [AutoLaunchedApplicationDictionary count]; i++)
//            {
//				if ([[[AutoLaunchedApplicationDictionary objectAtIndex:i] objectForKey:@"Path"] isEqualToString:[[NSBundle mainBundle] bundlePath]])
//                {
//					goto Path_already_there;
//				}
//			}
//			[AutoLaunchedApplicationDictionary addObject:sshAddDictionary];
//		}
//		
//		[loginDefaults setObject:AutoLaunchedApplicationDictionary forKey:@"AutoLaunchedApplicationDictionary"];
//		[defaults setPersistentDomain:loginDefaults forName:@"loginwindow"];
//	}
//	
//	[defaults synchronize];
//	
//  Path_already_there:

    setupLoginItem();

	NSString * dotMacOSXPath        = [[[NSString alloc] initWithFormat:@"%@/.MacOSX", NSHomeDirectoryForUser([NSString stringWithCString:logname])] autorelease];
	NSString * environmentPlistPath = [[[NSString alloc] initWithFormat:@"%@/.MacOSX/environment.plist", NSHomeDirectoryForUser([NSString stringWithCString:logname])] autorelease];
	NSMutableDictionary * environmentPlist;
	BOOL isDir;
	BOOL isExisting;
	NSFileManager * manager = [NSFileManager defaultManager];

	isExisting = [manager fileExistsAtPath:dotMacOSXPath isDirectory:&isDir];
	
	if (!isExisting)
    {
		if (![manager createDirectoryAtPath:dotMacOSXPath attributes:nil])
        {
            NSLog(@"failed to create directory %@", dotMacOSXPath);
			exit(EXIT_FAILURE);
		}
	}
	else if (isExisting && !isDir)
    {
		[manager movePath:dotMacOSXPath toPath:[NSString stringWithFormat:@"%@~", dotMacOSXPath] handler:nil];
		
		if (![manager createDirectoryAtPath:dotMacOSXPath attributes:nil])
        {
            NSLog(@"failed to create directory %@", dotMacOSXPath);
			exit(EXIT_FAILURE);
		}
	}

	if ([manager fileExistsAtPath:environmentPlistPath])
    {
		NSString * sshAgentPid;
		
		if (!(environmentPlist = [[NSMutableDictionary alloc] initWithContentsOfFile:environmentPlistPath]))
        {
			NSLog(@"failed to load %@.", environmentPlistPath);
			exit(EXIT_FAILURE);
		}
		
		[environmentPlist autorelease];
		
		if (sshAgentPid = [environmentPlist objectForKey:@"SSH_AGENT_PID"])
        {
			if (kill(atoi([sshAgentPid cString]), 0) == 0)
            {
                NSLog(@"ssh-agent is already running for user %s.", logname);
				exit(EXIT_SUCCESS);
			}
		}
	}
	else
    {
		environmentPlist = [[[NSMutableDictionary alloc] init] autorelease];
	}

	NSString *sshAskPass = [[[NSBundle mainBundle] executablePath] stringByAppendingString:@"_SSHAdd"];
	
	if (([[defaults objectForKey:@"enabled"] boolValue] == YES) || ![environmentPlist objectForKey:@"DISPLAY"])
    {
		[environmentPlist setObject:[defaults objectForKey:@"_DISPLAY"] forKey:@"DISPLAY"];
		setenv("DISPLAY", [[defaults objectForKey:@"_DISPLAY"] cString], 1);
	}
	
	if (([[defaults objectForKey:@"enabled"] boolValue] == YES) || ![environmentPlist objectForKey:@"SSH_ASKPASS"])
    {
		[environmentPlist setObject:sshAskPass forKey:@"SSH_ASKPASS"];
		setenv("SSH_ASKPASS", [sshAskPass cString], 1);
	}
	
	NSTask       * sshAgentTask = [[NSTask alloc] init];
	NSPipe       * sshAgentPipe = [NSPipe pipe];
	NSFileHandle * sshAgentPipeStdOut = [sshAgentPipe fileHandleForReading];
	NSData       * readData = nil;

	[sshAgentTask setStandardOutput:sshAgentPipe];
	[sshAgentTask setArguments:[NSArray arrayWithObjects:@"-s", nil]];
	[sshAgentTask setLaunchPath:@"/usr/bin/ssh-agent"];
	[sshAgentTask launch];
	
	readData = [sshAgentPipeStdOut readDataToEndOfFile];
	extractData(readData);
	
	[sshAgentTask release];
	
    NSLog(@"ssh-agent[%d]: %s started for user %s.", ssh_agent_pid, ssh_auth_sock, logname);
	
	[environmentPlist setObject:[NSString stringWithCString:ssh_auth_sock] forKey:@"SSH_AUTH_SOCK"];
	[environmentPlist setObject:[NSString stringWithFormat:@"%d", ssh_agent_pid] forKey:@"SSH_AGENT_PID"];
	
	if (![environmentPlist writeToFile:environmentPlistPath atomically:YES])
    {
		NSLog(@"failed to write %@.", environmentPlistPath);
		/* somethings wrong kill the agent.... ? */
		kill(ssh_agent_pid, SIGTERM);
		exit(EXIT_FAILURE);
	}
	
	return 0;
}

NSString * serviceString;

int main(int argc, const char * argv[])
{

	NSAutoreleasePool * pool     = [[NSAutoreleasePool alloc] init];
	NSUserDefaults    * defaults = [NSUserDefaults standardUserDefaults];
	
	/* need this in case System Prefs were never run */
	[defaults registerDefaults:[NSDictionary dictionaryWithObjectsAndKeys:
		[NSNumber numberWithBool:YES],					@"enabled",
		[NSNumber numberWithBool:NO],					@"UseUserSpecifiedKeyFile",
		[NSNumber numberWithBool:NO],					@"StorePassphraseInKeychain",
		[NSArray arrayWithObject:@"<default-ssh-key>"],	@"UserSpecifiedKeyFiles",
		nil]
	];
	
	serviceString = [[[[[[[NSProcessInfo processInfo]
		arguments]
		objectAtIndex:1]
		stringByTrimmingCharactersInSet:
			[NSCharacterSet whitespaceCharacterSet]]
		componentsSeparatedByString:@" "]
		lastObject]
		stringByTrimmingCharactersInSet:
			[NSCharacterSet characterSetWithCharactersInString:@":"]];
	
	NS_DURING
	
	if ((argc == 3) && [[[[NSProcessInfo processInfo] arguments] objectAtIndex:2] isEqualToString:@"-ssh-add"])
    {
		return NSApplicationMain(argc, argv);
	}
	else if ((argc == 3) && [[[[NSProcessInfo processInfo] arguments] objectAtIndex:2] isEqualToString:@"-LoginHook"])
    {
		struct passwd *pwd;
		
		if (!(pwd = getpwnam(argv[1])))
        {
            NSLog(@"ssh-add: getpwnam(%s) failed.", argv[1]);
			exit(EXIT_FAILURE);
		}
		
		return login_hook(argc, argv, pwd);
	}
	else
    {

		BOOL alreadyRanDefault = NO;
		NSMutableArray * sshAddArguments = [[[NSMutableArray alloc] init] autorelease];

		setenv("DISPLAY", ":0.0", 1);
		setenv("SSH_ASKPASS", [[[[NSBundle mainBundle] executablePath] stringByAppendingString:@"_SSHAdd"] cString], 1);
		
		if ([[defaults objectForKey:@"UseUserSpecifiedKeyFile"] boolValue] == YES)
        {
			for (unsigned int i = 0; i < [[defaults objectForKey:@"UserSpecifiedKeyFiles"] count]; i++)
            {
				if ([[[defaults objectForKey:@"UserSpecifiedKeyFiles"] objectAtIndex:i] isEqualToString:@"<default-ssh-key>"])
                {
					[NSTask launchedTaskWithLaunchPath:@"/usr/bin/ssh-add" arguments:[NSArray array]];
					alreadyRanDefault = YES;
				}
				else
                {
					[sshAddArguments addObject:[[defaults objectForKey:@"UserSpecifiedKeyFiles"] objectAtIndex:i]];
				}
			}
		}

		if (([sshAddArguments count] != 0) || !alreadyRanDefault)
        {
			[NSTask launchedTaskWithLaunchPath:@"/usr/bin/ssh-add" arguments:sshAddArguments];
		}

		exit(EXIT_SUCCESS);
	}

	NS_HANDLER

    NSLog(@"exception: %@", localException);

	NS_ENDHANDLER

	[pool release];

	return EXIT_FAILURE;

}